from langchain_core.documents import Document
from langchain_core.vectorstores import InMemoryVectorStore
from pydantic import BaseModel, Field
from typing import Optional, Any


class MyInMemoryVectorStore(BaseModel):
    embedding_model: Any = Field(..., description="Embedding model required.")
    vector_store: Any = Field(None, description="Vector store.")

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.vector_store = InMemoryVectorStore(embedding=self.embedding_model)

    async def add_documents(self, documents: list[Document], ids: list[str] | None = None) -> list[str]:
        document_ids = self.vector_store.add_documents(documents=documents, ids=ids)
        return document_ids

    async def del_documents(self, document_ids: list[str]) -> Optional[bool]:
        res = self.vector_store.delete(ids=document_ids)
        return res

    async def similarity_search(self, query: str, k: int = 4) -> list[Document]:
        res = self.vector_store.similarity_search(query=query, k=k)
        return res
